🚀 Deployed on https://63c72b8d51fc9a52f27f454d--nostalgic-wright-5fa355.netlify.app

Hi @Smit-create,

I have rewritten the exercise for the Kesten process in JAX. Could you have a quick look and let me know if there is anything I could change?

Many thanks.

Thanks for the quick update. This looks good to me!

Many thanks @Smit-create. I will pass this to @jstac for review.

Thanks @HumphreyYang @Smit-create

Comments:

• Why 4 in random.split(key, 4)?
• Does update_draws actually update T times into the future? I'm not clear on what scan is doing here.
• "Here is the solution in Jax" -> "Here is the solution in JAX" (perhaps link to Google JAX)
• Don't we need to install JAX via pip? (If so, perhaps at the start, with the other pip installs.) (CC @mmcky)
• "We can also use Numba with for loops" -> "We can ... for loops to generate the observations (replicating the results we obtained with JAX)"

Hi @jstac,

Thanks for the feedback. I have pushed an improved version based on your comments.

Does update_draws actually update T times into the future? I'm not clear on what scan is doing here.

I wrote a comparison of using for loop and using lax.scan to loop/scan through T periods in this notebook. The notebook will show that the results are the same, except the leading column is when $t=1$. What lax.scan does is it loops through parameters in the leading dimension, which is the dimension T, of (a_random, b_random, e_random). The function for lax.scan has two returns: one is the carrying state used as the input for the next stage, and a stacked matrix with all output concatenated together. In this case, the result from each state $t$ is carried through with res as the first return.

Don't we need to install JAX via pip? (If so, perhaps at the start, with the other pip installs.) (CC @mmcky)

I think pip for JAX has been included in the docker since the last update, so I did not use a pip install at the beginning.

Many thanks @HumphreyYang . Some comments:

1. I suggest as a convention that we use JAX's jit function as jax.jit rather than just jit to minimize confusion with jit from numba.
2. Regarding imports, I think our convention is to pip import anything not in the lastest version of Anaconda. @mmcky , can you confirm this point? If so, should we mention it in the manual?
3. Nice work with the comparison between the for loop and scan! I suggest that you include both, starting with the for loop and then moving to scan and explaining briefly how scan can be used to generate more efficient code.

However, regarding point 3, before you proceed, please check that the big timing difference persists when you run the functions a second time. I think a lot of that difference will be in compile time, since the compiler is unrolling the for loop.

CC @Smit-create regarding the conventions mentioned in my last comment.

However, regarding point 3, before you proceed, please check that the big timing difference persists when you run the functions a second time. I think a lot of that difference will be in compile time, since the compiler is unrolling the for loop.

It's an interesting experiment to run. The timing difference is still significant (around 900 ms vs 10 s) when the jitted function is run the second time. I think this might be due to how jit handles jax primitives and native Python code in compiled functions.

I will proceed to include both of them as a comparison, and I think it definitely helps when explaining what lax.scan is here for.

Thanks @HumphreyYang , very interesting. Please go ahead.

I also changed my code from using for loops to lax.scan which definitely is very very fast. Should we show this difference in each notebook? In my case, it was not even in seconds but went up to several minutes.

Hi @Smit-create , thanks for investigating. That's interesting. Was the acceleration in the compile time, execution time, or both?

Since this is a teaching resource, and since scan is a little unusual, I think it would be good to have the comparison between scan and the for loop in this case as well.

Was the acceleration in the compile time, execution time, or both?

I observed that the acceleration is mainly in the compile time (I tried several implementations and chose the fastest after changing some static args)

See the results:

%%time

w_jax_result = wealth_time_series_for_loop_jax(wdy.y_mean, ts_length, wdy, size)

Output:

CPU times: user 48.6 s, sys: 1.31 s, total: 50 s
Wall time: 50.5 s

2nd-time results:

CPU times: user 11.2 ms, sys: 0 ns, total: 11.2 ms
Wall time: 13.5 ms

While using lax.scan:
1st run:

CPU times: user 620 ms, sys: 9.76 ms, total: 630 ms
Wall time: 628 ms

2nd run:

CPU times: user 13.5 ms, sys: 1.03 ms, total: 14.5 ms
Wall time: 14.3 ms

Thanks @Smit-create , very interesting.

Compile time can be a real issue with JAX. I suggest that we include both for loop and scan code, with timing of both compile time and run time, along with some comments explaining why the timings are different.

The comments should avoid specific numbers, of course, since they will keep changing as hardware / software changes.

@mmcky , a friendly reminder that you are pinged above.

Please let me know when this is ready for my review.

Regarding imports, I think our convention is to pip import anything not in the lastest version of Anaconda. @mmcky , can you confirm this point? If so, should we mention it in the manual?

Thanks @jstac. This is a really interesting question. I would say this is true -- jax is a tricky one though given the hardware interaction. It is installed ahead of lecture builds to set it up correctly.

The one we need to use in lectures would be:

pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

it will also always need to be on its own line.

@HumphreyYang could we add this at the top of the lecture with any installs that are needed (that aren't included with anaconda)

Perhaps we need to write a blog post about installing jax to make use of gpu. Typically I have found when installing via pip without the specific drivers call from google it installs a cpu only version.

@HumphreyYang could we add this at the top of the lecture with any installs that are needed (that aren't included with anaconda)

Perhaps we need to write a blog post about installing jax to make use of gpu. Typically I have found when installing via pip without the specific drivers call from google it installs a cpu only version.

Thanks for the suggestion @mmcky. When adding the pip install, should I just add the lines below

!pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

but I reckon it will raise errors for readers without CUDA support or using windows machines and will not be necessary if readers are running the code on Colab. Therefore, my concern is that the errors raised by these two lines may confuse readers instead of helping them. Maybe we can include something like:

# If you are running on a machine without CUDA support, run the line below:
# !pip install --upgrade "jax[CPU]"

!pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

@HumphreyYang perhaps the best short term option (as you suggest) is to update the information note at the top of lectures that need gpu support

https://python-programming.quantecon.org/jax_intro.html#jax

and then update the note to include the items you noted above re: cpu and gpu options.

Our lectures actually won't need to install jax when building given the server is configured to have jax installed using cuda support already. so that should both help users and minimise additional installs.

Medium term I think we need to consider a page that discusses jax and the different options. I could work on that next week.

@HumphreyYang perhaps the best short term option (as you suggest) is to update the information note at the top of lectures that need gpu support

https://python-programming.quantecon.org/jax_intro.html#jax

and then update the note to include the items you noted above re: cpu and gpu options.

Our lectures actually won't need to install jax when building given the server is configured to have jax installed using cuda support already. so that should both help users and minimise additional installs.

Medium term I think we need to consider a page that discusses jax and the different options. I could work on that next week.

Just to make it clear as a short-run convention, we add this first

```{note}
This lecture is built using [hardware](status:machine-details) that
has access to a GPU. This means that 

1. the lecture might be significantly slower when running on your machine, and
2. the code is well-suited to execution with [Google colab](https://colab.research.google.com/github/QuantEcon/lecture-python-programming.notebooks/blob/master/jax_intro.ipynb)
```

and include

# If you are running on a machine without CUDA support, then run the line below:
# !pip install --upgrade "jax[CPU]"

# If you have CUDA support on your machine, then run the following:
# !pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

before imports. Is this what we should use for now, if I understand it correctly?

thanks @HumphreyYang just put the pip statements in code blocks 👍

Hi @jstac,

The GPU notes and import guide have been added. This is now ready for your review.

Many thanks.

@HumphreyYang rather than comments I think

# If you are running on a machine without CUDA support, then run the line below:
# !pip install --upgrade "jax[CPU]"

# If you have CUDA support on your machine, then run the following:
# !pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

should go in the note about using gpu hardware inside of code highlighted blocks.

Might make it too big but it would be great if you could add a screenshot to see what that would look like. We could also add a dropdown class to the note admonition (and a title to the note perhaps to dropdown the information for those that need it)

@HumphreyYang rather than comments I think

# If you are running on a machine without CUDA support, then run the line below:
# !pip install --upgrade "jax[CPU]"

# If you have CUDA support on your machine, then run the following:
# !pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

should go in the note about using gpu hardware inside of code hig
hlighted blocks.

Might make it too big but it would be great if you could add a screenshot to see what that would look like. We could also add a dropdown class to the note admonition (and a title to the note perhaps to dropdown the information for those that need it)

Hi @mmcky, I think the dropdown version will be better given the not is very large, but I think we need a better title for the note instead of "GPU Note" to attract readers to read this note:

instead of "GPU Note"

Something like "Install libraries based on GPU availability" would be nicer?

One concern is that every lecture should always run without any changes on a local machine with Anaconda installed.

(Otherwise we get a lot of complaints that our code is broken.)

Right now both JAX installs are commented out, so this condition fails.

I suggest that we do !pip install --upgrade jax[CPU] along with pip install quantecon etc. Then the condition stated above holds.

Second, in the GPU note --- which we should perhaps upgrade to a warning, we should offer two options:

1. Run on Colab and change the runtime type to add a GPU
2. Run locally and follow the pip install instructions for JAX with GPU support.

For 2, I suggest that we just make this a link rather than writing it out --- since the code won't be run on our server, it's not future proof.

I like @Smit-create 's suggestion for the title. We could also change it to "GPU Warning".

I think it's important that all readers read that note at least once.

Hi everyone,

Thanks for your feedback. I have pushed a new version with some changes to GPU Warning and pip command. Here is the latest deployment.

Please let me know if there is any change I could make.

Hi @jstac,

Our PRs are ready for review with GPU warnings updated. Please let us know if there is anything we need to change (CC @Smit-create)

Many thanks in advance.

This is very high quality. Well done @HumphreyYang (and thanks for reviewing @Smit-create ).

The use of scan really accelerates the code and it's a great example of how to use it.

With apologies, and thanks for your patience, a few more changes to the GPU notification please: (CC @Smit-create)

• Please have the warning open by default, rather than hidden.
• Please change "GPU Warning" to "GPU in use".
• "This lecture is built using" -> "This lecture is accelerated via"
• "and uses JAX" -> "and JAX"
• Cut sentence "As a result, the lecture will be... "

Regarding the code, please try to bring it back to around 80 characters per line, at least roughly. So please modify

    for t in range(T):
        s = s.at[:, t+1].set(jnp.where(s[:, t] < s_bar, 
                             jnp.exp(e_random[t, :]), 
                             jnp.exp(a_random[t, :]) * s[:, t] + jnp.exp(b_random[t, :])))

I suggest giving names or symbols to jnp.exp(a_random[t, :]) to shorten and bring the code a bit closer to the underlying maths.

This is looking good @HumphreyYang . Many thanks. Please ping me when the merge conflict is resolved and you are ready for me to merge.

Hi @jstac, the conflict is resolved. Please merge when you are available.

Many thanks @HumphreyYang , very good effort. Merging.